package top.houry.limit;

import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class TokenBucketRateLimiter {
    private final long capacity; // 桶的容量
    private final long tokens;   // 初始令牌数
    private final long refillRate; // 令牌填充速率（每秒添加的令牌数）
    private final long refillInterval; // 令牌填充间隔（毫秒）
    private long availableTokens; // 当前可用的令牌数
    private long lastRefillTimestamp; // 上次填充令牌的时间戳

    private final Lock lock = new ReentrantLock();

    public TokenBucketRateLimiter(long capacity, long tokens, long refillRate) {
        this.capacity = capacity;
        this.tokens = tokens;
        this.refillRate = refillRate;
        this.refillInterval = 1000 / refillRate; // 计算填充间隔
        this.availableTokens = tokens;
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    public boolean tryAcquire() {
        lock.lock();
        try {
            refillTokens();
            if (availableTokens >= 1) {
                availableTokens--;
                return true;
            } else {
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    private void refillTokens() {
        long now = System.currentTimeMillis();
        if (now - lastRefillTimestamp > refillInterval) {
            long tokensToAdd = (now - lastRefillTimestamp) / refillInterval * refillRate;
            availableTokens = Math.min(capacity, availableTokens + tokensToAdd);
            lastRefillTimestamp = now;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(100, 100, 100);

        for (int i = 0; i < 150; i++) {
            boolean result = rateLimiter.tryAcquire();
            System.out.println("请求 " + i + ": " + (result ? "通过" : "拒绝"));
//            if (i % 10 == 0) { // 模拟请求间隔
//                Thread.sleep(50); // 每10个请求暂停50毫秒
//            }
        }
    }
}
